{
 "cells": [
  {
   "cell_type": "raw",
   "id": "b61d713f-3ef8-4bef-8ff1-9859276c7726",
   "metadata": {},
   "source": [
    "[2024-09-26]\n",
    "pretrain 데이터 전처리 실행"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a969593-842a-469d-8ce8-c61518679c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import wfdb\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "_LEAD_NAMES = [\"I\", \"II\", \"III\", \"aVR\", \"aVL\", \"aVF\", \"V1\", \"V2\", \"V3\", \"V4\", \"V5\", \"V6\"]\n",
    "# _LEAD_NAMES = ['I', 'II', 'III', 'AVR', 'AVL', 'AVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']"
   ]
  },
  {
   "cell_type": "raw",
   "id": "fb898884-9912-4e53-b6f7-564962cbe1d8",
   "metadata": {},
   "source": [
    "if __name__ == \"__main__\":\n",
    "    run(get_parser())\n",
    "\n",
    "python process_ecg.py \\\n",
    "    --input_dir ${DATABASE_DIRECTORY} \\\n",
    "    --output_dir ${WAVEFORM_DIRECTORY} \\\n",
    "    --index_path ${INDEX_PATH}\n",
    "\n",
    "- input_dir: 원시 WFDB 파형 파일이 저장된 directory\n",
    "- output_dir: ECG PKL 파일 저장될 directory\n",
    "- index_path: index csv 파일 저장될 경로"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "319dfb1f-5ac8-46d3-8d3a-a9038939a83b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_parser():\n",
    "    description = \"Process WFDB ECG database.\"\n",
    "    # parser = argparse.ArgumentParser(description=description)\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('-i',\n",
    "                        '--input_dir',\n",
    "                        type=str,\n",
    "                        # required=True,\n",
    "                        default='/tf/physionet.org/files/challenge-2021/1.0.3/training/ningbo/',\n",
    "                        help=\"Path to the WFDB ECG database directory.\")\n",
    "    parser.add_argument('-o',\n",
    "                        '--output_dir',\n",
    "                        type=str,\n",
    "                        # required=True,\n",
    "                        default='./ningbo/ecgs/',\n",
    "                        help=\"Path to the directory where the preprocessed signals will be saved.\")\n",
    "    parser.add_argument('--index_path',\n",
    "                        type=str,\n",
    "                        default='./ningbo/index.csv',\n",
    "                        help=\"Path to the index file.\")\n",
    "    args = parser.parse_args(\"\")\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "41ae9fa0-ad53-4fa9-a5c6-e9163bddd7a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Namespace(input_dir='/tf/physionet.org/files/challenge-2021/1.0.3/training/ningbo/', output_dir='./ningbo/ecgs/', index_path='./ningbo/index.csv')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = get_parser()\n",
    "args"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e06139-70d8-4e27-a8d7-3ccf531599ee",
   "metadata": {},
   "source": [
    "get_parser() 부분 실행 완 => args 처리됨. 이후 run 코드 순서대로 실행하면 됨!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea547cef-aab4-487c-8911-56393c721ad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_records(root_dir):\n",
    "    \"\"\"Find all the .hea files in the root directory and its subdirectories.\n",
    "    Args:\n",
    "        root_dir (str): The directory to search for .hea files.\n",
    "    Returns:\n",
    "        records (set): A set of record names.\n",
    "                       (e.g., ['database/1/ecg001', 'database/1/ecg001', ..., 'database/9/ecg991'])\n",
    "    \"\"\"\n",
    "    records = set()\n",
    "    for root, _, files in os.walk(root_dir):\n",
    "        for file in files:\n",
    "            extension = os.path.splitext(file)[1]\n",
    "            if extension == '.hea':\n",
    "                record = os.path.relpath(os.path.join(root, file), root_dir)[:-4]\n",
    "                records.add(record)\n",
    "    records = sorted(records)\n",
    "    return records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a348247-59b8-4609-b0b9-c19bf09fd3f3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['g1/JS10647',\n",
       " 'g1/JS10648',\n",
       " 'g1/JS10649',\n",
       " 'g1/JS10650',\n",
       " 'g1/JS10651',\n",
       " 'g1/JS10652',\n",
       " 'g1/JS10653',\n",
       " 'g1/JS10654',\n",
       " 'g1/JS10655',\n",
       " 'g1/JS10656',\n",
       " 'g1/JS10657',\n",
       " 'g1/JS10658',\n",
       " 'g1/JS10659',\n",
       " 'g1/JS10660',\n",
       " 'g1/JS10661',\n",
       " 'g1/JS10662',\n",
       " 'g1/JS10663',\n",
       " 'g1/JS10664',\n",
       " 'g1/JS10665',\n",
       " 'g1/JS10666',\n",
       " 'g1/JS10667',\n",
       " 'g1/JS10668',\n",
       " 'g1/JS10669',\n",
       " 'g1/JS10670',\n",
       " 'g1/JS10671',\n",
       " 'g1/JS10672',\n",
       " 'g1/JS10673',\n",
       " 'g1/JS10674',\n",
       " 'g1/JS10675',\n",
       " 'g1/JS10676',\n",
       " 'g1/JS10677',\n",
       " 'g1/JS10678',\n",
       " 'g1/JS10679',\n",
       " 'g1/JS10680',\n",
       " 'g1/JS10681',\n",
       " 'g1/JS10682',\n",
       " 'g1/JS10683',\n",
       " 'g1/JS10684',\n",
       " 'g1/JS10685',\n",
       " 'g1/JS10686',\n",
       " 'g1/JS10687',\n",
       " 'g1/JS10688',\n",
       " 'g1/JS10689',\n",
       " 'g1/JS10690',\n",
       " 'g1/JS10691',\n",
       " 'g1/JS10692',\n",
       " 'g1/JS10693',\n",
       " 'g1/JS10694',\n",
       " 'g1/JS10695',\n",
       " 'g1/JS10696',\n",
       " 'g1/JS10697',\n",
       " 'g1/JS10698',\n",
       " 'g1/JS10699',\n",
       " 'g1/JS10700',\n",
       " 'g1/JS10701',\n",
       " 'g1/JS10702',\n",
       " 'g1/JS10703',\n",
       " 'g1/JS10704',\n",
       " 'g1/JS10705',\n",
       " 'g1/JS10706',\n",
       " 'g1/JS10707',\n",
       " 'g1/JS10708',\n",
       " 'g1/JS10709',\n",
       " 'g1/JS10710',\n",
       " 'g1/JS10711',\n",
       " 'g1/JS10712',\n",
       " 'g1/JS10713',\n",
       " 'g1/JS10714',\n",
       " 'g1/JS10715',\n",
       " 'g1/JS10716',\n",
       " 'g1/JS10717',\n",
       " 'g1/JS10718',\n",
       " 'g1/JS10719',\n",
       " 'g1/JS10720',\n",
       " 'g1/JS10721',\n",
       " 'g1/JS10722',\n",
       " 'g1/JS10723',\n",
       " 'g1/JS10724',\n",
       " 'g1/JS10725',\n",
       " 'g1/JS10726',\n",
       " 'g1/JS10727',\n",
       " 'g1/JS10728',\n",
       " 'g1/JS10729',\n",
       " 'g1/JS10730',\n",
       " 'g1/JS10731',\n",
       " 'g1/JS10732',\n",
       " 'g1/JS10733',\n",
       " 'g1/JS10734',\n",
       " 'g1/JS10735',\n",
       " 'g1/JS10736',\n",
       " 'g1/JS10737',\n",
       " 'g1/JS10738',\n",
       " 'g1/JS10739',\n",
       " 'g1/JS10740',\n",
       " 'g1/JS10741',\n",
       " 'g1/JS10742',\n",
       " 'g1/JS10743',\n",
       " 'g1/JS10744',\n",
       " 'g1/JS10745',\n",
       " 'g1/JS10746',\n",
       " 'g1/JS10747',\n",
       " 'g1/JS10748',\n",
       " 'g1/JS10749',\n",
       " 'g1/JS10750',\n",
       " 'g1/JS10751',\n",
       " 'g1/JS10752',\n",
       " 'g1/JS10753',\n",
       " 'g1/JS10754',\n",
       " 'g1/JS10755',\n",
       " 'g1/JS10756',\n",
       " 'g1/JS10757',\n",
       " 'g1/JS10758',\n",
       " 'g1/JS10759',\n",
       " 'g1/JS10760',\n",
       " 'g1/JS10761',\n",
       " 'g1/JS10762',\n",
       " 'g1/JS10763',\n",
       " 'g1/JS10764',\n",
       " 'g1/JS10765',\n",
       " 'g1/JS10766',\n",
       " 'g1/JS10767',\n",
       " 'g1/JS10768',\n",
       " 'g1/JS10769',\n",
       " 'g1/JS10770',\n",
       " 'g1/JS10771',\n",
       " 'g1/JS10772',\n",
       " 'g1/JS10773',\n",
       " 'g1/JS10774',\n",
       " 'g1/JS10775',\n",
       " 'g1/JS10776',\n",
       " 'g1/JS10777',\n",
       " 'g1/JS10778',\n",
       " 'g1/JS10779',\n",
       " 'g1/JS10780',\n",
       " 'g1/JS10781',\n",
       " 'g1/JS10782',\n",
       " 'g1/JS10783',\n",
       " 'g1/JS10784',\n",
       " 'g1/JS10785',\n",
       " 'g1/JS10786',\n",
       " 'g1/JS10787',\n",
       " 'g1/JS10788',\n",
       " 'g1/JS10789',\n",
       " 'g1/JS10790',\n",
       " 'g1/JS10791',\n",
       " 'g1/JS10792',\n",
       " 'g1/JS10793',\n",
       " 'g1/JS10794',\n",
       " 'g1/JS10795',\n",
       " 'g1/JS10796',\n",
       " 'g1/JS10797',\n",
       " 'g1/JS10798',\n",
       " 'g1/JS10799',\n",
       " 'g1/JS10800',\n",
       " 'g1/JS10801',\n",
       " 'g1/JS10802',\n",
       " 'g1/JS10803',\n",
       " 'g1/JS10804',\n",
       " 'g1/JS10805',\n",
       " 'g1/JS10806',\n",
       " 'g1/JS10807',\n",
       " 'g1/JS10808',\n",
       " 'g1/JS10809',\n",
       " 'g1/JS10810',\n",
       " 'g1/JS10811',\n",
       " 'g1/JS10812',\n",
       " 'g1/JS10813',\n",
       " 'g1/JS10814',\n",
       " 'g1/JS10815',\n",
       " 'g1/JS10816',\n",
       " 'g1/JS10817',\n",
       " 'g1/JS10818',\n",
       " 'g1/JS10819',\n",
       " 'g1/JS10820',\n",
       " 'g1/JS10821',\n",
       " 'g1/JS10822',\n",
       " 'g1/JS10823',\n",
       " 'g1/JS10824',\n",
       " 'g1/JS10825',\n",
       " 'g1/JS10826',\n",
       " 'g1/JS10827',\n",
       " 'g1/JS10828',\n",
       " 'g1/JS10829',\n",
       " 'g1/JS10830',\n",
       " 'g1/JS10831',\n",
       " 'g1/JS10832',\n",
       " 'g1/JS10833',\n",
       " 'g1/JS10834',\n",
       " 'g1/JS10835',\n",
       " 'g1/JS10836',\n",
       " 'g1/JS10837',\n",
       " 'g1/JS10838',\n",
       " 'g1/JS10839',\n",
       " 'g1/JS10840',\n",
       " 'g1/JS10841',\n",
       " 'g1/JS10842',\n",
       " 'g1/JS10843',\n",
       " 'g1/JS10844',\n",
       " 'g1/JS10845',\n",
       " 'g1/JS10846',\n",
       " 'g1/JS10847',\n",
       " 'g1/JS10848',\n",
       " 'g1/JS10849',\n",
       " 'g1/JS10850',\n",
       " 'g1/JS10851',\n",
       " 'g1/JS10852',\n",
       " 'g1/JS10853',\n",
       " 'g1/JS10854',\n",
       " 'g1/JS10855',\n",
       " 'g1/JS10856',\n",
       " 'g1/JS10857',\n",
       " 'g1/JS10858',\n",
       " 'g1/JS10859',\n",
       " 'g1/JS10860',\n",
       " 'g1/JS10861',\n",
       " 'g1/JS10862',\n",
       " 'g1/JS10863',\n",
       " 'g1/JS10864',\n",
       " 'g1/JS10865',\n",
       " 'g1/JS10866',\n",
       " 'g1/JS10867',\n",
       " 'g1/JS10868',\n",
       " 'g1/JS10869',\n",
       " 'g1/JS10870',\n",
       " 'g1/JS10871',\n",
       " 'g1/JS10872',\n",
       " 'g1/JS10873',\n",
       " 'g1/JS10874',\n",
       " 'g1/JS10875',\n",
       " 'g1/JS10876',\n",
       " 'g1/JS10877',\n",
       " 'g1/JS10878',\n",
       " 'g1/JS10879',\n",
       " 'g1/JS10880',\n",
       " 'g1/JS10881',\n",
       " 'g1/JS10882',\n",
       " 'g1/JS10883',\n",
       " 'g1/JS10884',\n",
       " 'g1/JS10885',\n",
       " 'g1/JS10886',\n",
       " 'g1/JS10887',\n",
       " 'g1/JS10888',\n",
       " 'g1/JS10889',\n",
       " 'g1/JS10890',\n",
       " 'g1/JS10891',\n",
       " 'g1/JS10892',\n",
       " 'g1/JS10893',\n",
       " 'g1/JS10894',\n",
       " 'g1/JS10895',\n",
       " 'g1/JS10896',\n",
       " 'g1/JS10897',\n",
       " 'g1/JS10898',\n",
       " 'g1/JS10899',\n",
       " 'g1/JS10900',\n",
       " 'g1/JS10901',\n",
       " 'g1/JS10902',\n",
       " 'g1/JS10903',\n",
       " 'g1/JS10904',\n",
       " 'g1/JS10905',\n",
       " 'g1/JS10906',\n",
       " 'g1/JS10907',\n",
       " 'g1/JS10908',\n",
       " 'g1/JS10909',\n",
       " 'g1/JS10910',\n",
       " 'g1/JS10911',\n",
       " 'g1/JS10912',\n",
       " 'g1/JS10913',\n",
       " 'g1/JS10914',\n",
       " 'g1/JS10915',\n",
       " 'g1/JS10916',\n",
       " 'g1/JS10917',\n",
       " 'g1/JS10918',\n",
       " 'g1/JS10919',\n",
       " 'g1/JS10920',\n",
       " 'g1/JS10921',\n",
       " 'g1/JS10922',\n",
       " 'g1/JS10923',\n",
       " 'g1/JS10924',\n",
       " 'g1/JS10925',\n",
       " 'g1/JS10926',\n",
       " 'g1/JS10927',\n",
       " 'g1/JS10928',\n",
       " 'g1/JS10929',\n",
       " 'g1/JS10930',\n",
       " 'g1/JS10931',\n",
       " 'g1/JS10932',\n",
       " 'g1/JS10933',\n",
       " 'g1/JS10934',\n",
       " 'g1/JS10935',\n",
       " 'g1/JS10936',\n",
       " 'g1/JS10937',\n",
       " 'g1/JS10938',\n",
       " 'g1/JS10939',\n",
       " 'g1/JS10940',\n",
       " 'g1/JS10941',\n",
       " 'g1/JS10942',\n",
       " 'g1/JS10943',\n",
       " 'g1/JS10944',\n",
       " 'g1/JS10945',\n",
       " 'g1/JS10946',\n",
       " 'g1/JS10947',\n",
       " 'g1/JS10948',\n",
       " 'g1/JS10949',\n",
       " 'g1/JS10950',\n",
       " 'g1/JS10951',\n",
       " 'g1/JS10952',\n",
       " 'g1/JS10953',\n",
       " 'g1/JS10954',\n",
       " 'g1/JS10955',\n",
       " 'g1/JS10956',\n",
       " 'g1/JS10957',\n",
       " 'g1/JS10958',\n",
       " 'g1/JS10959',\n",
       " 'g1/JS10960',\n",
       " 'g1/JS10961',\n",
       " 'g1/JS10962',\n",
       " 'g1/JS10963',\n",
       " 'g1/JS10964',\n",
       " 'g1/JS10965',\n",
       " 'g1/JS10966',\n",
       " 'g1/JS10967',\n",
       " 'g1/JS10968',\n",
       " 'g1/JS10969',\n",
       " 'g1/JS10970',\n",
       " 'g1/JS10971',\n",
       " 'g1/JS10972',\n",
       " 'g1/JS10973',\n",
       " 'g1/JS10974',\n",
       " 'g1/JS10975',\n",
       " 'g1/JS10976',\n",
       " 'g1/JS10977',\n",
       " 'g1/JS10978',\n",
       " 'g1/JS10979',\n",
       " 'g1/JS10980',\n",
       " 'g1/JS10981',\n",
       " 'g1/JS10982',\n",
       " 'g1/JS10983',\n",
       " 'g1/JS10984',\n",
       " 'g1/JS10985',\n",
       " 'g1/JS10986',\n",
       " 'g1/JS10987',\n",
       " 'g1/JS10988',\n",
       " 'g1/JS10989',\n",
       " 'g1/JS10990',\n",
       " 'g1/JS10991',\n",
       " 'g1/JS10992',\n",
       " 'g1/JS10993',\n",
       " 'g1/JS10994',\n",
       " 'g1/JS10995',\n",
       " 'g1/JS10996',\n",
       " 'g1/JS10997',\n",
       " 'g1/JS10998',\n",
       " 'g1/JS10999',\n",
       " 'g1/JS11000',\n",
       " 'g1/JS11001',\n",
       " 'g1/JS11002',\n",
       " 'g1/JS11003',\n",
       " 'g1/JS11004',\n",
       " 'g1/JS11005',\n",
       " 'g1/JS11006',\n",
       " 'g1/JS11007',\n",
       " 'g1/JS11008',\n",
       " 'g1/JS11009',\n",
       " 'g1/JS11010',\n",
       " 'g1/JS11011',\n",
       " 'g1/JS11012',\n",
       " 'g1/JS11013',\n",
       " 'g1/JS11014',\n",
       " 'g1/JS11015',\n",
       " 'g1/JS11016',\n",
       " 'g1/JS11017',\n",
       " 'g1/JS11018',\n",
       " 'g1/JS11019',\n",
       " 'g1/JS11020',\n",
       " 'g1/JS11021',\n",
       " 'g1/JS11022',\n",
       " 'g1/JS11023',\n",
       " 'g1/JS11024',\n",
       " 'g1/JS11025',\n",
       " 'g1/JS11026',\n",
       " 'g1/JS11027',\n",
       " 'g1/JS11028',\n",
       " 'g1/JS11029',\n",
       " 'g1/JS11030',\n",
       " 'g1/JS11031',\n",
       " 'g1/JS11032',\n",
       " 'g1/JS11033',\n",
       " 'g1/JS11034',\n",
       " 'g1/JS11035',\n",
       " 'g1/JS11036',\n",
       " 'g1/JS11037',\n",
       " 'g1/JS11038',\n",
       " 'g1/JS11039',\n",
       " 'g1/JS11040',\n",
       " 'g1/JS11041',\n",
       " 'g1/JS11042',\n",
       " 'g1/JS11043',\n",
       " 'g1/JS11044',\n",
       " 'g1/JS11045',\n",
       " 'g1/JS11046',\n",
       " 'g1/JS11047',\n",
       " 'g1/JS11048',\n",
       " 'g1/JS11049',\n",
       " 'g1/JS11050',\n",
       " 'g1/JS11051',\n",
       " 'g1/JS11052',\n",
       " 'g1/JS11053',\n",
       " 'g1/JS11054',\n",
       " 'g1/JS11055',\n",
       " 'g1/JS11056',\n",
       " 'g1/JS11057',\n",
       " 'g1/JS11058',\n",
       " 'g1/JS11059',\n",
       " 'g1/JS11060',\n",
       " 'g1/JS11061',\n",
       " 'g1/JS11062',\n",
       " 'g1/JS11063',\n",
       " 'g1/JS11064',\n",
       " 'g1/JS11065',\n",
       " 'g1/JS11066',\n",
       " 'g1/JS11067',\n",
       " 'g1/JS11068',\n",
       " 'g1/JS11069',\n",
       " 'g1/JS11070',\n",
       " 'g1/JS11071',\n",
       " 'g1/JS11072',\n",
       " 'g1/JS11073',\n",
       " 'g1/JS11074',\n",
       " 'g1/JS11075',\n",
       " 'g1/JS11076',\n",
       " 'g1/JS11077',\n",
       " 'g1/JS11078',\n",
       " 'g1/JS11079',\n",
       " 'g1/JS11080',\n",
       " 'g1/JS11081',\n",
       " 'g1/JS11082',\n",
       " 'g1/JS11083',\n",
       " 'g1/JS11084',\n",
       " 'g1/JS11085',\n",
       " 'g1/JS11086',\n",
       " 'g1/JS11087',\n",
       " 'g1/JS11088',\n",
       " 'g1/JS11089',\n",
       " 'g1/JS11090',\n",
       " 'g1/JS11091',\n",
       " 'g1/JS11092',\n",
       " 'g1/JS11093',\n",
       " 'g1/JS11094',\n",
       " 'g1/JS11095',\n",
       " 'g1/JS11096',\n",
       " 'g1/JS11097',\n",
       " 'g1/JS11098',\n",
       " 'g1/JS11099',\n",
       " 'g1/JS11100',\n",
       " 'g1/JS11101',\n",
       " 'g1/JS11102',\n",
       " 'g1/JS11103',\n",
       " 'g1/JS11104',\n",
       " 'g1/JS11105',\n",
       " 'g1/JS11106',\n",
       " 'g1/JS11107',\n",
       " 'g1/JS11108',\n",
       " 'g1/JS11109',\n",
       " 'g1/JS11110',\n",
       " 'g1/JS11111',\n",
       " 'g1/JS11112',\n",
       " 'g1/JS11113',\n",
       " 'g1/JS11114',\n",
       " 'g1/JS11115',\n",
       " 'g1/JS11116',\n",
       " 'g1/JS11117',\n",
       " 'g1/JS11118',\n",
       " 'g1/JS11119',\n",
       " 'g1/JS11120',\n",
       " 'g1/JS11121',\n",
       " 'g1/JS11122',\n",
       " 'g1/JS11123',\n",
       " 'g1/JS11124',\n",
       " 'g1/JS11125',\n",
       " 'g1/JS11126',\n",
       " 'g1/JS11127',\n",
       " 'g1/JS11128',\n",
       " 'g1/JS11129',\n",
       " 'g1/JS11130',\n",
       " 'g1/JS11131',\n",
       " 'g1/JS11132',\n",
       " 'g1/JS11133',\n",
       " 'g1/JS11134',\n",
       " 'g1/JS11135',\n",
       " 'g1/JS11136',\n",
       " 'g1/JS11137',\n",
       " 'g1/JS11138',\n",
       " 'g1/JS11139',\n",
       " 'g1/JS11140',\n",
       " 'g1/JS11141',\n",
       " 'g1/JS11142',\n",
       " 'g1/JS11143',\n",
       " 'g1/JS11144',\n",
       " 'g1/JS11145',\n",
       " 'g1/JS11146',\n",
       " 'g1/JS11147',\n",
       " 'g1/JS11148',\n",
       " 'g1/JS11149',\n",
       " 'g1/JS11150',\n",
       " 'g1/JS11151',\n",
       " 'g1/JS11152',\n",
       " 'g1/JS11153',\n",
       " 'g1/JS11154',\n",
       " 'g1/JS11155',\n",
       " 'g1/JS11156',\n",
       " 'g1/JS11157',\n",
       " 'g1/JS11158',\n",
       " 'g1/JS11159',\n",
       " 'g1/JS11160',\n",
       " 'g1/JS11161',\n",
       " 'g1/JS11162',\n",
       " 'g1/JS11163',\n",
       " 'g1/JS11164',\n",
       " 'g1/JS11165',\n",
       " 'g1/JS11166',\n",
       " 'g1/JS11167',\n",
       " 'g1/JS11168',\n",
       " 'g1/JS11169',\n",
       " 'g1/JS11170',\n",
       " 'g1/JS11171',\n",
       " 'g1/JS11172',\n",
       " 'g1/JS11173',\n",
       " 'g1/JS11174',\n",
       " 'g1/JS11175',\n",
       " 'g1/JS11176',\n",
       " 'g1/JS11177',\n",
       " 'g1/JS11178',\n",
       " 'g1/JS11179',\n",
       " 'g1/JS11180',\n",
       " 'g1/JS11181',\n",
       " 'g1/JS11182',\n",
       " 'g1/JS11183',\n",
       " 'g1/JS11184',\n",
       " 'g1/JS11185',\n",
       " 'g1/JS11186',\n",
       " 'g1/JS11187',\n",
       " 'g1/JS11188',\n",
       " 'g1/JS11189',\n",
       " 'g1/JS11190',\n",
       " 'g1/JS11191',\n",
       " 'g1/JS11192',\n",
       " 'g1/JS11193',\n",
       " 'g1/JS11194',\n",
       " 'g1/JS11195',\n",
       " 'g1/JS11196',\n",
       " 'g1/JS11197',\n",
       " 'g1/JS11198',\n",
       " 'g1/JS11199',\n",
       " 'g1/JS11200',\n",
       " 'g1/JS11201',\n",
       " 'g1/JS11202',\n",
       " 'g1/JS11203',\n",
       " 'g1/JS11204',\n",
       " 'g1/JS11205',\n",
       " 'g1/JS11206',\n",
       " 'g1/JS11207',\n",
       " 'g1/JS11208',\n",
       " 'g1/JS11209',\n",
       " 'g1/JS11210',\n",
       " 'g1/JS11211',\n",
       " 'g1/JS11212',\n",
       " 'g1/JS11213',\n",
       " 'g1/JS11214',\n",
       " 'g1/JS11215',\n",
       " 'g1/JS11216',\n",
       " 'g1/JS11217',\n",
       " 'g1/JS11218',\n",
       " 'g1/JS11219',\n",
       " 'g1/JS11220',\n",
       " 'g1/JS11221',\n",
       " 'g1/JS11222',\n",
       " 'g1/JS11223',\n",
       " 'g1/JS11224',\n",
       " 'g1/JS11225',\n",
       " 'g1/JS11226',\n",
       " 'g1/JS11227',\n",
       " 'g1/JS11228',\n",
       " 'g1/JS11229',\n",
       " 'g1/JS11230',\n",
       " 'g1/JS11231',\n",
       " 'g1/JS11232',\n",
       " 'g1/JS11233',\n",
       " 'g1/JS11234',\n",
       " 'g1/JS11235',\n",
       " 'g1/JS11236',\n",
       " 'g1/JS11237',\n",
       " 'g1/JS11238',\n",
       " 'g1/JS11239',\n",
       " 'g1/JS11240',\n",
       " 'g1/JS11241',\n",
       " 'g1/JS11242',\n",
       " 'g1/JS11243',\n",
       " 'g1/JS11244',\n",
       " 'g1/JS11245',\n",
       " 'g1/JS11246',\n",
       " 'g1/JS11247',\n",
       " 'g1/JS11248',\n",
       " 'g1/JS11249',\n",
       " 'g1/JS11250',\n",
       " 'g1/JS11251',\n",
       " 'g1/JS11252',\n",
       " 'g1/JS11253',\n",
       " 'g1/JS11254',\n",
       " 'g1/JS11255',\n",
       " 'g1/JS11256',\n",
       " 'g1/JS11257',\n",
       " 'g1/JS11258',\n",
       " 'g1/JS11259',\n",
       " 'g1/JS11260',\n",
       " 'g1/JS11261',\n",
       " 'g1/JS11262',\n",
       " 'g1/JS11263',\n",
       " 'g1/JS11264',\n",
       " 'g1/JS11265',\n",
       " 'g1/JS11266',\n",
       " 'g1/JS11267',\n",
       " 'g1/JS11268',\n",
       " 'g1/JS11269',\n",
       " 'g1/JS11270',\n",
       " 'g1/JS11271',\n",
       " 'g1/JS11272',\n",
       " 'g1/JS11273',\n",
       " 'g1/JS11274',\n",
       " 'g1/JS11275',\n",
       " 'g1/JS11276',\n",
       " 'g1/JS11277',\n",
       " 'g1/JS11278',\n",
       " 'g1/JS11279',\n",
       " 'g1/JS11280',\n",
       " 'g1/JS11281',\n",
       " 'g1/JS11282',\n",
       " 'g1/JS11283',\n",
       " 'g1/JS11284',\n",
       " 'g1/JS11285',\n",
       " 'g1/JS11286',\n",
       " 'g1/JS11287',\n",
       " 'g1/JS11288',\n",
       " 'g1/JS11289',\n",
       " 'g1/JS11290',\n",
       " 'g1/JS11291',\n",
       " 'g1/JS11292',\n",
       " 'g1/JS11293',\n",
       " 'g1/JS11294',\n",
       " 'g1/JS11295',\n",
       " 'g1/JS11296',\n",
       " 'g1/JS11297',\n",
       " 'g1/JS11298',\n",
       " 'g1/JS11299',\n",
       " 'g1/JS11300',\n",
       " 'g1/JS11301',\n",
       " 'g1/JS11302',\n",
       " 'g1/JS11303',\n",
       " 'g1/JS11304',\n",
       " 'g1/JS11305',\n",
       " 'g1/JS11306',\n",
       " 'g1/JS11307',\n",
       " 'g1/JS11308',\n",
       " 'g1/JS11309',\n",
       " 'g1/JS11310',\n",
       " 'g1/JS11311',\n",
       " 'g1/JS11312',\n",
       " 'g1/JS11313',\n",
       " 'g1/JS11314',\n",
       " 'g1/JS11315',\n",
       " 'g1/JS11316',\n",
       " 'g1/JS11317',\n",
       " 'g1/JS11318',\n",
       " 'g1/JS11319',\n",
       " 'g1/JS11320',\n",
       " 'g1/JS11321',\n",
       " 'g1/JS11322',\n",
       " 'g1/JS11323',\n",
       " 'g1/JS11324',\n",
       " 'g1/JS11325',\n",
       " 'g1/JS11326',\n",
       " 'g1/JS11327',\n",
       " 'g1/JS11328',\n",
       " 'g1/JS11329',\n",
       " 'g1/JS11330',\n",
       " 'g1/JS11331',\n",
       " 'g1/JS11332',\n",
       " 'g1/JS11333',\n",
       " 'g1/JS11334',\n",
       " 'g1/JS11335',\n",
       " 'g1/JS11336',\n",
       " 'g1/JS11337',\n",
       " 'g1/JS11338',\n",
       " 'g1/JS11339',\n",
       " 'g1/JS11340',\n",
       " 'g1/JS11341',\n",
       " 'g1/JS11342',\n",
       " 'g1/JS11343',\n",
       " 'g1/JS11344',\n",
       " 'g1/JS11345',\n",
       " 'g1/JS11346',\n",
       " 'g1/JS11347',\n",
       " 'g1/JS11348',\n",
       " 'g1/JS11349',\n",
       " 'g1/JS11350',\n",
       " 'g1/JS11351',\n",
       " 'g1/JS11352',\n",
       " 'g1/JS11353',\n",
       " 'g1/JS11354',\n",
       " 'g1/JS11355',\n",
       " 'g1/JS11356',\n",
       " 'g1/JS11357',\n",
       " 'g1/JS11358',\n",
       " 'g1/JS11359',\n",
       " 'g1/JS11360',\n",
       " 'g1/JS11361',\n",
       " 'g1/JS11362',\n",
       " 'g1/JS11363',\n",
       " 'g1/JS11364',\n",
       " 'g1/JS11365',\n",
       " 'g1/JS11366',\n",
       " 'g1/JS11367',\n",
       " 'g1/JS11368',\n",
       " 'g1/JS11369',\n",
       " 'g1/JS11370',\n",
       " 'g1/JS11371',\n",
       " 'g1/JS11372',\n",
       " 'g1/JS11373',\n",
       " 'g1/JS11374',\n",
       " 'g1/JS11375',\n",
       " 'g1/JS11376',\n",
       " 'g1/JS11377',\n",
       " 'g1/JS11378',\n",
       " 'g1/JS11379',\n",
       " 'g1/JS11380',\n",
       " 'g1/JS11381',\n",
       " 'g1/JS11382',\n",
       " 'g1/JS11383',\n",
       " 'g1/JS11384',\n",
       " 'g1/JS11385',\n",
       " 'g1/JS11386',\n",
       " 'g1/JS11387',\n",
       " 'g1/JS11388',\n",
       " 'g1/JS11389',\n",
       " 'g1/JS11390',\n",
       " 'g1/JS11391',\n",
       " 'g1/JS11392',\n",
       " 'g1/JS11393',\n",
       " 'g1/JS11394',\n",
       " 'g1/JS11395',\n",
       " 'g1/JS11396',\n",
       " 'g1/JS11397',\n",
       " 'g1/JS11398',\n",
       " 'g1/JS11399',\n",
       " 'g1/JS11400',\n",
       " 'g1/JS11401',\n",
       " 'g1/JS11402',\n",
       " 'g1/JS11403',\n",
       " 'g1/JS11404',\n",
       " 'g1/JS11405',\n",
       " 'g1/JS11406',\n",
       " 'g1/JS11407',\n",
       " 'g1/JS11408',\n",
       " 'g1/JS11409',\n",
       " 'g1/JS11410',\n",
       " 'g1/JS11411',\n",
       " 'g1/JS11412',\n",
       " 'g1/JS11413',\n",
       " 'g1/JS11414',\n",
       " 'g1/JS11415',\n",
       " 'g1/JS11416',\n",
       " 'g1/JS11417',\n",
       " 'g1/JS11418',\n",
       " 'g1/JS11419',\n",
       " 'g1/JS11420',\n",
       " 'g1/JS11421',\n",
       " 'g1/JS11422',\n",
       " 'g1/JS11423',\n",
       " 'g1/JS11424',\n",
       " 'g1/JS11425',\n",
       " 'g1/JS11426',\n",
       " 'g1/JS11427',\n",
       " 'g1/JS11428',\n",
       " 'g1/JS11429',\n",
       " 'g1/JS11430',\n",
       " 'g1/JS11431',\n",
       " 'g1/JS11432',\n",
       " 'g1/JS11433',\n",
       " 'g1/JS11434',\n",
       " 'g1/JS11435',\n",
       " 'g1/JS11436',\n",
       " 'g1/JS11437',\n",
       " 'g1/JS11438',\n",
       " 'g1/JS11439',\n",
       " 'g1/JS11440',\n",
       " 'g1/JS11441',\n",
       " 'g1/JS11442',\n",
       " 'g1/JS11443',\n",
       " 'g1/JS11444',\n",
       " 'g1/JS11445',\n",
       " 'g1/JS11446',\n",
       " 'g1/JS11447',\n",
       " 'g1/JS11448',\n",
       " 'g1/JS11449',\n",
       " 'g1/JS11450',\n",
       " 'g1/JS11451',\n",
       " 'g1/JS11452',\n",
       " 'g1/JS11453',\n",
       " 'g1/JS11454',\n",
       " 'g1/JS11455',\n",
       " 'g1/JS11456',\n",
       " 'g1/JS11457',\n",
       " 'g1/JS11458',\n",
       " 'g1/JS11459',\n",
       " 'g1/JS11460',\n",
       " 'g1/JS11461',\n",
       " 'g1/JS11462',\n",
       " 'g1/JS11463',\n",
       " 'g1/JS11464',\n",
       " 'g1/JS11465',\n",
       " 'g1/JS11466',\n",
       " 'g1/JS11467',\n",
       " 'g1/JS11468',\n",
       " 'g1/JS11469',\n",
       " 'g1/JS11470',\n",
       " 'g1/JS11471',\n",
       " 'g1/JS11472',\n",
       " 'g1/JS11473',\n",
       " 'g1/JS11474',\n",
       " 'g1/JS11475',\n",
       " 'g1/JS11476',\n",
       " 'g1/JS11477',\n",
       " 'g1/JS11478',\n",
       " 'g1/JS11479',\n",
       " 'g1/JS11480',\n",
       " 'g1/JS11481',\n",
       " 'g1/JS11482',\n",
       " 'g1/JS11483',\n",
       " 'g1/JS11484',\n",
       " 'g1/JS11485',\n",
       " 'g1/JS11486',\n",
       " 'g1/JS11487',\n",
       " 'g1/JS11488',\n",
       " 'g1/JS11489',\n",
       " 'g1/JS11490',\n",
       " 'g1/JS11491',\n",
       " 'g1/JS11492',\n",
       " 'g1/JS11493',\n",
       " 'g1/JS11494',\n",
       " 'g1/JS11495',\n",
       " 'g1/JS11496',\n",
       " 'g1/JS11497',\n",
       " 'g1/JS11498',\n",
       " 'g1/JS11499',\n",
       " 'g1/JS11500',\n",
       " 'g1/JS11501',\n",
       " 'g1/JS11502',\n",
       " 'g1/JS11503',\n",
       " 'g1/JS11504',\n",
       " 'g1/JS11505',\n",
       " 'g1/JS11506',\n",
       " 'g1/JS11507',\n",
       " 'g1/JS11508',\n",
       " 'g1/JS11509',\n",
       " 'g1/JS11510',\n",
       " 'g1/JS11511',\n",
       " 'g1/JS11512',\n",
       " 'g1/JS11513',\n",
       " 'g1/JS11514',\n",
       " 'g1/JS11515',\n",
       " 'g1/JS11516',\n",
       " 'g1/JS11517',\n",
       " 'g1/JS11518',\n",
       " 'g1/JS11519',\n",
       " 'g1/JS11520',\n",
       " 'g1/JS11521',\n",
       " 'g1/JS11522',\n",
       " 'g1/JS11523',\n",
       " 'g1/JS11524',\n",
       " 'g1/JS11525',\n",
       " 'g1/JS11526',\n",
       " 'g1/JS11527',\n",
       " 'g1/JS11528',\n",
       " 'g1/JS11529',\n",
       " 'g1/JS11530',\n",
       " 'g1/JS11531',\n",
       " 'g1/JS11532',\n",
       " 'g1/JS11533',\n",
       " 'g1/JS11534',\n",
       " 'g1/JS11535',\n",
       " 'g1/JS11536',\n",
       " 'g1/JS11537',\n",
       " 'g1/JS11538',\n",
       " 'g1/JS11539',\n",
       " 'g1/JS11540',\n",
       " 'g1/JS11541',\n",
       " 'g1/JS11542',\n",
       " 'g1/JS11543',\n",
       " 'g1/JS11544',\n",
       " 'g1/JS11545',\n",
       " 'g1/JS11546',\n",
       " 'g1/JS11547',\n",
       " 'g1/JS11548',\n",
       " 'g1/JS11549',\n",
       " 'g1/JS11550',\n",
       " 'g1/JS11551',\n",
       " 'g1/JS11552',\n",
       " 'g1/JS11553',\n",
       " 'g1/JS11554',\n",
       " 'g1/JS11555',\n",
       " 'g1/JS11556',\n",
       " 'g1/JS11557',\n",
       " 'g1/JS11558',\n",
       " 'g1/JS11559',\n",
       " 'g1/JS11560',\n",
       " 'g1/JS11561',\n",
       " 'g1/JS11562',\n",
       " 'g1/JS11563',\n",
       " 'g1/JS11564',\n",
       " 'g1/JS11565',\n",
       " 'g1/JS11566',\n",
       " 'g1/JS11567',\n",
       " 'g1/JS11568',\n",
       " 'g1/JS11569',\n",
       " 'g1/JS11570',\n",
       " 'g1/JS11571',\n",
       " 'g1/JS11572',\n",
       " 'g1/JS11573',\n",
       " 'g1/JS11574',\n",
       " 'g1/JS11575',\n",
       " 'g1/JS11576',\n",
       " 'g1/JS11577',\n",
       " 'g1/JS11578',\n",
       " 'g1/JS11579',\n",
       " 'g1/JS11580',\n",
       " 'g1/JS11581',\n",
       " 'g1/JS11582',\n",
       " 'g1/JS11583',\n",
       " 'g1/JS11584',\n",
       " 'g1/JS11585',\n",
       " 'g1/JS11586',\n",
       " 'g1/JS11587',\n",
       " 'g1/JS11588',\n",
       " 'g1/JS11589',\n",
       " 'g1/JS11590',\n",
       " 'g1/JS11591',\n",
       " 'g1/JS11592',\n",
       " 'g1/JS11593',\n",
       " 'g1/JS11594',\n",
       " 'g1/JS11595',\n",
       " 'g1/JS11596',\n",
       " 'g1/JS11597',\n",
       " 'g1/JS11598',\n",
       " 'g1/JS11599',\n",
       " 'g1/JS11600',\n",
       " 'g1/JS11601',\n",
       " 'g1/JS11602',\n",
       " 'g1/JS11603',\n",
       " 'g1/JS11604',\n",
       " 'g1/JS11605',\n",
       " 'g1/JS11606',\n",
       " 'g1/JS11607',\n",
       " 'g1/JS11608',\n",
       " 'g1/JS11609',\n",
       " 'g1/JS11610',\n",
       " 'g1/JS11611',\n",
       " 'g1/JS11612',\n",
       " 'g1/JS11613',\n",
       " 'g1/JS11614',\n",
       " 'g1/JS11615',\n",
       " 'g1/JS11616',\n",
       " 'g1/JS11617',\n",
       " 'g1/JS11618',\n",
       " 'g1/JS11619',\n",
       " 'g1/JS11620',\n",
       " 'g1/JS11621',\n",
       " 'g1/JS11622',\n",
       " 'g1/JS11623',\n",
       " 'g1/JS11624',\n",
       " 'g1/JS11625',\n",
       " 'g1/JS11626',\n",
       " 'g1/JS11627',\n",
       " 'g1/JS11628',\n",
       " 'g1/JS11629',\n",
       " 'g1/JS11630',\n",
       " 'g1/JS11631',\n",
       " 'g1/JS11632',\n",
       " 'g1/JS11633',\n",
       " 'g1/JS11634',\n",
       " 'g1/JS11635',\n",
       " 'g1/JS11636',\n",
       " 'g1/JS11637',\n",
       " 'g1/JS11638',\n",
       " 'g1/JS11639',\n",
       " 'g1/JS11640',\n",
       " 'g1/JS11641',\n",
       " 'g1/JS11642',\n",
       " 'g1/JS11643',\n",
       " 'g1/JS11644',\n",
       " 'g1/JS11645',\n",
       " 'g10/JS19646',\n",
       " ...]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Identify the header fiels\n",
    "record_rel_paths = find_records(args.input_dir)\n",
    "record_rel_paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f7d2f1b1-edbb-41e5-a649-e567581fe63c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(34905,\n",
       " ['g1/JS10647', 'g1/JS10648', 'g1/JS10649', 'g1/JS10650'],\n",
       " ['g9/JS19642', 'g9/JS19643', 'g9/JS19644', 'g9/JS19645'])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(record_rel_paths), record_rel_paths[:4], record_rel_paths[-4:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b5851faa-c9ae-400b-bff6-04a61dd07294",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 34905 records.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Found {len(record_rel_paths)} records.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0d5e3c20-bf34-4f6e-b09e-7b80d868ed7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RELATIVE_FILE_PATH</th>\n",
       "      <th>FILE_NAME</th>\n",
       "      <th>SAMPLE_RATE</th>\n",
       "      <th>SOURCE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: [RELATIVE_FILE_PATH, FILE_NAME, SAMPLE_RATE, SOURCE]\n",
       "Index: []"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Prepare an index dataframe\n",
    "index_df = pd.DataFrame(columns = [\"RELATIVE_FILE_PATH\", \"FILE_NAME\", \"SAMPLE_RATE\", \"SOURCE\"])\n",
    "index_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e672e6b6-eafd-4ab2-809e-3b00523dce2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def moving_window_crop(x: np.ndarray, crop_length: int, crop_stride: int) -> np.ndarray:\n",
    "    \"\"\"Crop the input sequence with a moving window.\n",
    "    \"\"\"\n",
    "    if crop_length > x.shape[1]:\n",
    "        raise ValueError(f\"crop_length must be smaller than the length of x ({x.shape[1]}).\")\n",
    "    start_idx = np.arange(0, x.shape[1] - crop_length + 1, crop_stride)\n",
    "    return [x[:, i:i + crop_length] for i in start_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5f8a34a1-f355-49da-a8b3-c38887eeda2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████| 34905/34905 [2:21:30<00:00,  4.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved 34808 cropped signals.\n"
     ]
    }
   ],
   "source": [
    "num_saved = 0\n",
    "for record_rel_path in tqdm(record_rel_paths):\n",
    "    record_rel_dir, record_name = os.path.split(record_rel_path)\n",
    "    save_dir = os.path.join(args.output_dir, record_rel_dir)\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    source_name = record_rel_dir.split(\"/\")[0]\n",
    "    signal, record_info = wfdb.rdsamp(os.path.join(args.input_dir, record_rel_path))\n",
    "    lead_idx = np.array([record_info[\"sig_name\"].index(lead_name) for lead_name in _LEAD_NAMES])\n",
    "    signal = signal[:, lead_idx]\n",
    "    fs = record_info[\"fs\"]\n",
    "    signal_length = record_info[\"sig_len\"]\n",
    "    if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "        continue\n",
    "    cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "    for idx, cropped_signal in enumerate(cropped_signals):\n",
    "        if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "            continue\n",
    "        pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "                     os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "        index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "                                   f\"{record_name}_{idx}.pkl\",\n",
    "                                   fs,\n",
    "                                   source_name]\n",
    "        num_saved += 1\n",
    "\n",
    "print(f\"Saved {num_saved} cropped signals.\")\n",
    "os.makedirs(os.path.dirname(args.index_path), exist_ok=True)\n",
    "index_df.to_csv(args.index_path, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bf8b378c-f587-4d73-b10b-8c9e747c68da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RELATIVE_FILE_PATH</th>\n",
       "      <th>FILE_NAME</th>\n",
       "      <th>SAMPLE_RATE</th>\n",
       "      <th>SOURCE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>g1/JS10647_0.pkl</td>\n",
       "      <td>JS10647_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>g1/JS10648_0.pkl</td>\n",
       "      <td>JS10648_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>g1/JS10649_0.pkl</td>\n",
       "      <td>JS10649_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>g1/JS10650_0.pkl</td>\n",
       "      <td>JS10650_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>g1/JS10651_0.pkl</td>\n",
       "      <td>JS10651_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34803</th>\n",
       "      <td>g9/JS19641_0.pkl</td>\n",
       "      <td>JS19641_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34804</th>\n",
       "      <td>g9/JS19642_0.pkl</td>\n",
       "      <td>JS19642_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34805</th>\n",
       "      <td>g9/JS19643_0.pkl</td>\n",
       "      <td>JS19643_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34806</th>\n",
       "      <td>g9/JS19644_0.pkl</td>\n",
       "      <td>JS19644_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34807</th>\n",
       "      <td>g9/JS19645_0.pkl</td>\n",
       "      <td>JS19645_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>34808 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      RELATIVE_FILE_PATH      FILE_NAME  SAMPLE_RATE SOURCE\n",
       "0       g1/JS10647_0.pkl  JS10647_0.pkl          500     g1\n",
       "1       g1/JS10648_0.pkl  JS10648_0.pkl          500     g1\n",
       "2       g1/JS10649_0.pkl  JS10649_0.pkl          500     g1\n",
       "3       g1/JS10650_0.pkl  JS10650_0.pkl          500     g1\n",
       "4       g1/JS10651_0.pkl  JS10651_0.pkl          500     g1\n",
       "...                  ...            ...          ...    ...\n",
       "34803   g9/JS19641_0.pkl  JS19641_0.pkl          500     g9\n",
       "34804   g9/JS19642_0.pkl  JS19642_0.pkl          500     g9\n",
       "34805   g9/JS19643_0.pkl  JS19643_0.pkl          500     g9\n",
       "34806   g9/JS19644_0.pkl  JS19644_0.pkl          500     g9\n",
       "34807   g9/JS19645_0.pkl  JS19645_0.pkl          500     g9\n",
       "\n",
       "[34808 rows x 4 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628c586f-b93e-40de-a931-61003405b633",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d2a9b9-27e5-4454-b245-b5d9fdd6c2a2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a9fd448-58b2-48c5-86e0-c2d4f632d2db",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b6aab8-9d33-4d4c-85dc-4123ccf7ff7a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0d95ab-d658-47d6-a21b-8693bccfa614",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44de267-612f-49f0-a0c4-014f0a3c9066",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "145e3b32-4569-4098-b8f9-05967f08d266",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bed473c-f05b-40af-a2d5-320b009589f6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a7fe86-4137-437a-a1f5-4c6d781daef6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0751a74c-96c9-4f1c-b3d5-000636565096",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf5d001-0cc0-4772-b232-7f84b35995ec",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d00647b-9456-4fe1-b189-a10633064f56",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f93b79-336d-43cb-9b07-9075e8c9d0c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save all the cropped signals\n",
    "num_saved = 0\n",
    "for record_rel_path in tqdm(record_rel_paths):\n",
    "    record_rel_dir, record_name = os.path.split(record_rel_path)\n",
    "    save_dir = os.path.join(args.output_dir, record_rel_dir)\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    source_name = record_rel_dir.split(\"/\")[0]\n",
    "    signal, record_info = wfdb.rdsamp(os.path.join(args.input_dir, record_rel_path))\n",
    "    lead_idx = np.array([record_info[\"sig_name\"].index(lead_name) for lead_name in _LEAD_NAMES])\n",
    "    signal = signal[:, lead_idx]\n",
    "    fs = record_info[\"fs\"]\n",
    "    signal_length = record_info[\"sig_len\"]\n",
    "    # if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "    #     continue\n",
    "    # cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "    # for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    #     if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "    #         continue\n",
    "    #     pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "    #                  os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "    #     index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "    #                                f\"{record_name}_{idx}.pkl\",\n",
    "    #                                fs,\n",
    "    #                                source_name]\n",
    "    #     num_saved += 1\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acc7bac7-3a8b-485f-a8ae-66f7cecb3ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "record_rel_dir, record_name, save_dir, source_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1ee7f9b-5f69-493a-b3ed-2d11ec010d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67d5325a-e0d3-43d2-a882-9d00cbb02b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "record_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff5b601-6f4d-4004-8b15-162aa5675b4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "    continue\n",
    "cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "        continue\n",
    "    pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "                 os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "    index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "                               f\"{record_name}_{idx}.pkl\",\n",
    "                               fs,\n",
    "                               source_name]\n",
    "    num_saved += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1521670-34dd-4150-9d70-2d43c9369593",
   "metadata": {},
   "outputs": [],
   "source": [
    "if signal_length < 10 * fs:\n",
    "    print(1)\n",
    "else: print(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd94d4d2-a3ef-4966-9c63-0ee79b8ca72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5580e89a-0341-424a-95cc-639a29be1a6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs * 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cacee1e-f867-487e-b215-be5f364aa365",
   "metadata": {},
   "outputs": [],
   "source": [
    "cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "cropped_signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c86396f1-2cd3-4056-b0a0-9a1bd81042bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(signal).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "123cec3f-1d7f-4073-85e5-d03c8dfddf8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(cropped_signals).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bc5c577-af73-435f-b278-b79b700baaed",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    print(cropped_signal)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15bd445-ac06-4e61-a41b-4ffc91df9395",
   "metadata": {},
   "outputs": [],
   "source": [
    "[f\"{record_rel_path}_{idx}.pkl\", f\"{record_name}_{idx}.pkl\", fs, source_name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81139e44-8652-4d9a-83e9-0dabab4f0c6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = pd.read_csv('/tf/physionet.org/files/ptb-xl/1.0.3/ptbxl_database.csv', index_col='ecg_id')\n",
    "Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efc3e165-3730-4b01-b336-8be6a3246514",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y.scp_codes # = Y.scp_codes.apply(lambda x: ast.literal_eval(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6c68783-152a-4e85-8cc6-6288905f0b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ast\n",
    "Y.scp_codes = Y.scp_codes.apply(lambda x: ast.literal_eval(x))\n",
    "Y.scp_codes"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
